import os
import json
from paddlex import create_pipeline
from collections import defaultdict


def process_seal_recognition(input_dir="cropped_seals", output_dir="output/seal_recognition")->dict:
    """
    处理裁剪后的印章图片，进行文字识别并输出结构化结果

    参数:
        input_dir: 输入目录(包含裁剪后的印章图片)
        output_dir: 输出目录(保存识别结果)
    """
    print("开始执行seal_rec 产线")
    # 创建输出目录
    os.makedirs(output_dir, exist_ok=True)

    # 创建识别流水线
    pipeline = create_pipeline(pipeline="seal_recognition")

    # 存储最终结果
    # final_results = defaultdict(lambda: defaultdict(dict))

    result_map = defaultdict(lambda: defaultdict(list))

    # 获取所有图片文件
    image_files = [f for f in os.listdir(input_dir) if f.lower().endswith(('.png', '.jpg', '.jpeg'))]

    for img_file in image_files:
        # 解析文件名 (格式: "图片名称-文字类型-数字类型-result中该类别的id")
        parts = os.path.splitext(img_file)[0].split('-')
        if len(parts) < 4:
            print(f"跳过文件名格式不正确的文件: {img_file}")
            continue

        original_name = parts[0]  # 原始图片名称
        text_type = parts[-3]  # 文字类型
        num_type = parts[-2]  # 数字类型
        class_id = parts[-1]  # 类别ID

        # 完整图片路径
        img_path = os.path.join(input_dir, img_file)

        # 进行印章识别
        try:
            output = pipeline.predict(
                img_path,
                use_doc_orientation_classify=False,
                use_doc_unwarping=False,
            )

            # 处理识别结果
            for res in output:

                seal_res_list = res.get('seal_res_list', [])
                if seal_res_list:
                    rec_text = seal_res_list[0].get('rec_texts', [])
                    rec_scores = seal_res_list[0].get('rec_scores', [])
                    key=f"{original_name}{os.path.splitext(img_file)[1]}"
                    result_map[key][num_type].append( {class_id: {"rec_text": rec_text, "rec_scores": rec_scores}})

                # 保存可视化结果和JSON(可选)
                res.save_to_img(output_dir)
                res.save_to_json(output_dir)

        except Exception as e:
            print(f"处理文件 {img_file} 时出错: {str(e)}")
            continue

    # 转换为普通字典并保存最终结果
    # final_results = {k: dict(v) for k, v in final_results.items()}
    # output_file = os.path.join(output_dir, "seal_recognition_results.json")

    # with open(output_file, 'w', encoding='utf-8') as f:
    #     json.dump(final_results, f, ensure_ascii=False, indent=2)

    print(f"识别完成，结果已保存到: {output_dir}")
    # return final_results
    return result_map



if __name__ == "__main__":
    # 使用示例
    results = process_seal_recognition(
        input_dir="cropped_seals",
        output_dir="output/seal_recognition"
    )


